package org.example.demo.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.example.demo.service.AgentAssistant;
import org.example.demo.service.CalculatorTools;
import org.example.demo.service.KnowledgeBaseAssistant;
import org.example.demo.service.WeatherTools;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class LangChain4jConfig {

    @Value("${langchain4j.open-ai.chat-model.api-key}")
    private String apiKey;

    @Value("${langchain4j.open-ai.chat-model.base-url}")
    private String baseUrl;

    @Value("${langchain4j.open-ai.chat-model.model-name}")
    private String modelName;

    @Value("${langchain4j.open-ai.embedding-model.model-name}")
    private String embeddingModelName;

    @Value("${langchain4j.open-ai.chat-model.temperature}")
    private Double temperature;

    @Value("${langchain4j.open-ai.chat-model.max-tokens}")
    private Integer maxTokens;

    @Bean("openAiChatModel")
    public ChatModel openAiChatModel() {
        return OpenAiChatModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .modelName(modelName)
                .temperature(temperature)
                .maxTokens(maxTokens)
                .build();
    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        // 这个Bean会为每个用户（由memoryId标识）创建一个MessageWindowChatMemory实例
        // withMaxMessages(10)表示每个对话最多保留最近的10条消息
        return memoryId -> MessageWindowChatMemory.withMaxMessages(10);
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        // 通常嵌入模型也使用相同的api-key和base-url
        // 注意：OpenAI有专门的嵌入模型名称，
        // 比如 "text-embedding-ada-002" 。
        // 如果不指定，LangChain4j可能会使用一个默认值。
        // 为清晰起见，最好在properties中也定义它。
        return OpenAiEmbeddingModel.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                // 推荐在application.properties中添加:
                // langchain4j.open-ai.embedding-model.model-name=text-embedding-ada-002
                .modelName(embeddingModelName)
                .build();
    }


    // 步骤1: 创建ContentRetriever Bean
    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore, EmbeddingModel embeddingModel) {
        // 这个检索器知道如何从向量存储中根据语义相似度检索内容
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3) // 每次检索最多返回3个最相关的片段
                .build();
    }

    // 步骤2: 创建RetrievalAugmentor Bean
    @Bean
    public RetrievalAugmentor retrievalAugmentor(ContentRetriever contentRetriever) {
        // 这是默认的检索增强器，它会使用提供的内容检索器
        return DefaultRetrievalAugmentor.builder()
                .contentRetriever(contentRetriever)
                .build();
    }

    // 步骤3: 创建最终的、具备RAG能力的AI服务Bean
    @Bean
    public KnowledgeBaseAssistant knowledgeBaseAssistant(ChatModel chatLanguageModel, RetrievalAugmentor retrievalAugmentor) {
        return AiServices.builder(KnowledgeBaseAssistant.class)
                .chatModel(chatLanguageModel)
                .retrievalAugmentor(retrievalAugmentor)
                .build();
    }

    // 修改Agent服务接口Bean
    @Bean
    public AgentAssistant agentAssistant(ChatModel chatLanguageModel,
                                         CalculatorTools calculatorTools,
                                         WeatherTools weatherTools) { // 注入新工具
        return AiServices.builder(AgentAssistant.class)
                .chatModel(chatLanguageModel)
                .tools(calculatorTools, weatherTools) // <-- 关键步骤：将所有工具都注册进去
                .build();
    }

}